'''
@Company: TWL
@Author: xue jian
@Email: xuejian@kanzhun.com
@Date: 2020-07-12 14:39:00
'''
'''
5211. 概率最大的路径
给你一个由 n 个节点（下标从 0 开始）组成的无向加权图，该图由一个描述边的列表组成，其中 edges[i] = [a, b] 表示连接节点 a 和 b 的一条无向边，且该边遍历成功的概率为 succProb[i] 。

指定两个节点分别作为起点 start 和终点 end ，请你找出从起点到终点成功概率最大的路径，并返回其成功概率。

如果不存在从 start 到 end 的路径，请 返回 0 。只要答案与标准答案的误差不超过 1e-5 ，就会被视作正确答案。

 

示例 1：



输入：n = 3, edges = [[0,1],[1,2],[0,2]], succProb = [0.5,0.5,0.2], start = 0, end = 2
输出：0.25000
解释：从起点到终点有两条路径，其中一条的成功概率为 0.2 ，而另一条为 0.5 * 0.5 = 0.25
示例 2：



输入：n = 3, edges = [[0,1],[1,2],[0,2]], succProb = [0.5,0.5,0.3], start = 0, end = 2
输出：0.30000
示例 3：



输入：n = 3, edges = [[0,1]], succProb = [0.5], start = 0, end = 2
输出：0.00000
解释：节点 0 和 节点 2 之间不存在路径
 

提示：

2 <= n <= 10^4
0 <= start, end < n
start != end
0 <= a, b < n
a != b
0 <= succProb.length == edges.length <= 2*10^4
0 <= succProb[i] <= 1
每两个节点之间最多有一条边
'''
'''
tips:将概率转成权重w = -log(p)，这个题变成最短路径问题。直接上Dijkstra。这是无向图，注意处理即可。
'''
from typing import List
class Solution:
    def maxProbability(self, n: int, edges: List[List[int]], succProb: List[float], start: int, end: int) -> float:
        import math
        import sys
        # print(math.log(0.2))
        dis = {}
        for i, v in enumerate(edges):
            # print(i[0], i[1])
            if v[0] not in dis:
                dis[v[0]] = {}
            if v[1] not in dis:
                dis[v[1]] = {}
            dis[v[0]][v[1]] = -1*math.log(succProb[i])
            dis[v[1]][v[0]] = -1*math.log(succProb[i])
        path = set()
        path.add(start)
        path_v = [start]
        re = 0
        tmp_start = start
        if start not in dis:
            return 0
        
        while True:
            tmp_dis = sys.maxsize
            tmp_v = -1
            for k,v in dis[start].items():
                if v < tmp_dis and k not in path:
                    # print('hha', k, exp(-v))
                    tmp_dis = v
                    tmp_v = k
            if tmp_v == -1:
                return 0
            # print(tmp_v, exp(-tmp_dis), dis)
            re = tmp_dis
            path_v.append(tmp_v)
            path.add(tmp_v)
            if tmp_v == end:
                # print(path_v, re)
                return math.exp(-re)
            for k, v in dis[tmp_v].items():
                if k in dis[start]:
                    dis[start][k] = min(dis[start][k], dis[start][tmp_v]+v)
                else:
                    dis[start][k] = dis[start][tmp_v]+v
            del(dis[start][tmp_v])
        return 0
            
if __name__ == "__main__":
    solution = Solution()
    n = 3
    edges = [[0,1],[1,2],[0,2]]
    succProb = [0.5,0.5,0.2]
    start = 0
    end = 2
    print(solution.maxProbability(n, edges, succProb, start, end))